from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from prompt.system_prompt import SYSTEM_PROMPT
from langchain_classic.memory import ConversationSummaryMemory, ConversationBufferWindowMemory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableSequence
from user_profile import UserProfile

chat_model = ChatOpenAI(
    model="Qwen3-8B:Q5_K_M",
    streaming=True,
    base_url="http://localhost:11434/v1",  # 关键：Ollama 暴露的 OpenAI 兼容端点
    api_key="ollama",
    temperature=0.7,
    timeout=30,  # 可选：设置超时时间（秒），避免长时间无响应
    max_completion_tokens=20 * 1024
)

user_exp_role = UserProfile(chat_model)

# ---------------------- 2. 初始化两个核心记忆组件 ----------------------
# ① 窗口缓冲记忆：保留最近 3 轮完整对话（细节不丢失）
window_memory = ConversationBufferWindowMemory(
    k=15,  # 关键参数：保留最近 3 轮（1 轮 = 用户+AI，共 6 条消息）
    return_messages=True,  # 返回 Message 对象，适配 Prompt
    input_key="user_input",
    output_key="ai_output"
)

# ---------------------- 3. 自定义记忆组合逻辑（核心）----------------------
def load_combined_memory(user_input):
    # 加载两个记忆的内容
    recent_dialog = window_memory.load_memory_variables({})["history"]  # 近期 3 轮完整对话

    return {
        "history": recent_dialog,
        "user_input": user_input
    }

# ---------------------- 4. 构建 Prompt + 对话链 ----------------------
prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    MessagesPlaceholder(variable_name="history"),  # 动态注入组合记忆
    ("human", f"记忆摘要如下：{user_exp_role.summarizes}, 你需要根据记忆摘要与用户进行交流，不要编造和脑补，不要过多提问。" + "用户最新输入：{user_input}")
])

# 串联：记忆加载 → Prompt 渲染 → LLM 生成
chain = RunnableSequence(
    load_combined_memory,
    prompt,
    chat_model
)


# ---------------------- 5. 多轮对话 + 流式输出 + 记忆同步 ----------------------
def multi_turn_chat():
    while True:
        user_input = input("用户：").strip()
        if user_input in ["退出", "quit", "exit"]:
            user_exp_role.save_state()
            break
        if not user_input:
            continue
        
        user_exp_role.summary(window_memory, user_input)

        full_reply = []

        try:
            # 流式迭代 Token，逐字输出
            for chunk in chain.stream(user_input):
                token = chunk.content if chunk.content else ""
                if token:
                    print(token, end="", flush=True)
                    full_reply.append(token)
            assistant_reply = "".join(full_reply).strip()
            print("\n")

            # 同步当前对话到两个记忆组件（关键：确保下次加载时包含最新对话）
            window_memory.save_context({"user_input": user_input}, {"ai_output": assistant_reply})

        except Exception as e:
            print(f"\n晚星：哎呀～ 好像出了点小问题😣，换个话题聊聊吧！")
            print(f"调试信息：{str(e)[:100]}\n")
            continue

if __name__ == "__main__":
    multi_turn_chat()
